# Create server parameters for stdio connection
from mcp import ClientSession, StdioServerParameters
from mcp.client.stdio import stdio_client
from langchain_mcp_adapters.tools import load_mcp_tools
from langgraph.prebuilt import create_react_agent
from langchain_openai import ChatOpenAI
import asyncio

import dotenv
dotenv.load_dotenv()

model = ChatOpenAI(model="gpt-4o-mini-2024-07-18")

server_params = StdioServerParameters(
    command="python",    
    # Make sure to update to the full absolute path to your math_server.py file    
    args=["server.py"],
)

async def run_agent():
    async with stdio_client(server_params) as (read, write):
        async with ClientSession(read, write) as session:       
            # Initialize the connection            
            await session.initialize()            
            
            # Get tools            
            tools = await load_mcp_tools(session)            
            
            # Create and run the agent            
            agent = create_react_agent(model, tools)            
            agent_response = await agent.ainvoke({"messages": "请根据你现有的tools介绍一下你能做什么？"})            
            return agent_response
            
# Run the async function
if __name__ == "__main__":
    result = asyncio.run(run_agent())    
    print(result['messages'][-1].content)